#!/usr/bin/env python

from xml.dom.minidom import parse
from shutil import copyfile
import datetime
import sys
import os

# XML utils
def getText(parentNode):
	rc = []
	for node in parentNode.childNodes:
		if node.nodeType == node.TEXT_NODE:
			rc.append(node.data)
	return ''.join(rc)

def getFirstElementChild(parentNode, name):
	for node in parentNode.childNodes:
		if node.nodeType == node.ELEMENT_NODE:
			if node.nodeName == name:
				return node
	return None

def countChildElements(parentNode):
	count = 0
	for node in parentNode.childNodes:
		if node.nodeType == node.ELEMENT_NODE:
			count = count + 1
	return count

def printChildElements(parentNode):
	for node in parentNode.childNodes:
		if node.nodeType == node.ELEMENT_NODE:
			print node.nodeName

def getNextElementSibling(node):
	node = node.nextSibling
	while node is not None and node.nodeType != node.ELEMENT_NODE: 
		node = node.nextSibling
	return node

# AlbumData.xml utils
def printAlbumChildElements(parentNode):
	for node in parentNode.childNodes:
		if node.nodeType == node.ELEMENT_NODE:
			print node.nodeName + " " ,
			if node.nodeName == 'key' or node.nodeName == 'string' or node.nodeName == 'integer':
				print "'" + getText(node) + "'"
			elif node.nodeName == 'array' or node.nodeName == 'dict':
				print "(" + str(countChildElements(node)) + ")"
			else:
				print

def getKeyedElement(parentNode, name):
	for node in parentNode.childNodes:
		if node.nodeType == node.ELEMENT_NODE:
			if node.nodeName == 'key':
				if getText(node) == name:
					return getNextElementSibling(node)

def getAppleTime(value):
    return datetime.datetime.fromtimestamp(978307200 + float(value))

def countPictures(rollList):
	count = 0
	node = getFirstElementChild(rollList,'dict')
	while node != None:
		count = count + int(getText(getKeyedElement(node, 'PhotoCount')))
		node = getNextElementSibling(node)
	return count

def processAlbum(album, mil):
	global destdir
	year = getAppleTime(getText(getKeyedElement(album, 'RollDateAsTimerInterval'))).year
	path = os.path.join(destdir, str(year))
	if not os.path.exists(path):
		os.mkdir(path)
	albumName = getText(getKeyedElement(album, 'RollName'))
	path = getFileName(path, albumName)
	os.mkdir(path)
	imgkey = getFirstElementChild(getKeyedElement(album, 'KeyList'),'string')
	while imgkey is not None:
		img = getKeyedElement(mil, getText(imgkey))
		processImage(img, path)
		imgkey = getNextElementSibling(imgkey)

def getFileName (path, name):
	fname = os.path.join(path, name)
	root, ext = os.path.splitext(name)
	count = 1
	while os.path.exists(fname):
		fname = os.path.join(path, root + '_' + str(count) + ext)
		count = count + 1
	return fname

def processImage(img, destPath):
	global imgdone
	global imgcount
	imgpath = getText(getKeyedElement(img, 'ImagePath'))
	org = getKeyedElement(img, 'OriginalPath')
	# If you want to copy the modified files, comment these 2 lines
	if org is not None:
		imgpath = getText(org)
	fname = getFileName(destPath, os.path.split(imgpath)[1])
	copyfile(imgpath, fname)
	imgdone = imgdone + 1
	printProgressBar (imgdone, imgcount)

# Other utils
def printProgressBar (done, total):
	if (done == total):
		print "All Done!"
	elif (done % 5 == 0):
		prog = int(float(done)/total * 50)
		text = "[" + ("#" * prog) + (" " * (50 - prog)) + "] " + str(done) + "/" + str(total)
		sys.stdout.write(text)
		sys.stdout.flush()
		sys.stdout.write("\b" * len(text))

def debug():
	global destdir
	global imgcount
	global imgdone
	global dom
	global mil
	global album
	
	dom = parse('/Users/rsc/Pictures/iPhoto Library/AlbumData.xml')
	destdir = '/Users/rsc/Downloads/testdir'
	dic = getFirstElementChild(getFirstElementChild(dom, 'plist'), 'dict')
	rollList = getKeyedElement(dic, 'List of Rolls')
	imgcount = 127
	imgdone = 0
	mil = getKeyedElement(dic, 'Master Image List')
	album = getFirstElementChild(rollList,'dict')

# Main program
def main():
	global destdir
	global imgcount
	global imgdone
	imgdone = 0
	if len(sys.argv) < 2:
		sys.stderr.write('Usage: sys.argv[0] AlbumData.xml DestDirectory')
		sys.exit(1)

	print "Reading AlbumData (this may take a bit)"
	dom = parse(sys.argv[1])
	destdir = sys.argv[2]
	dic = getFirstElementChild(getFirstElementChild(dom, 'plist'), 'dict')
	rollList = getKeyedElement(dic, 'List of Rolls')
	imgcount = countPictures(rollList)
	imgdone = 0
	print "Processing " + str(imgcount) + " images"
	mil = getKeyedElement(dic, 'Master Image List')

	node = getFirstElementChild(rollList,'dict')
	while node != None:
		processAlbum(node, mil)
		node = getNextElementSibling(node)

if __name__ == "__main__":
    main()
